"""
自动生成 smi 记录文件
"""

import os
from pathlib import Path


class SmiRecord:
    path_this = Path(__file__).resolve()
    flag_npu = 0
    flag_gpu = 1

    def __init__(self, name, device_type: int, record_path: Path):
        """_summary_
        Args:
            record_path (Path): path to save the log file
        """
        self.name = name
        self.device_type = device_type
        self.record_path = record_path

    def generate_simple_implement(self, interval=60):
        """Simplest implement of smi record, just echo the data to the file every 60s.

        Args:
            interval (_type_): default 60s
        """
        path_bash = Path(f"/tmp/smi_record/{self.name}/${{RANK}}_smi.sh")
        str_pre_env = "\n\n###### START generate_smi_shell ######\n"
        if self.device_type == self.flag_npu:
            str_cmd_smi = "npu-smi info"
        elif self.device_type == self.flag_gpu:
            str_cmd_smi = "nvidia-smi"

        str_cmd = f"""
mkdir -p {path_bash.parent}
mkdir -p {self.record_path.parent}
path_smi_shell={path_bash}
echo "
hostname -I >> {self.record_path}
ifconfig >> {self.record_path}
while true
do
    {str_cmd_smi} >> {self.record_path}
    date >> {self.record_path}
    sleep {interval}
done
" > ${{path_smi_shell}}
chmod +x ${{path_smi_shell}}
{path_bash} &

smi_pid=$(jobs -l | grep {path_bash.parent} | awk '{{print $2}}')
"""
        str_cmd += "echo 'smi record started.'"
        str_post_env = "\n###### END generate_smi_shell ######\n"
        return str_pre_env + str_cmd + str_post_env


if __name__ == "__main__":
    pass
